from sqlalchemy import create_engine, text
from sqlalchemy.engine import Engine
from datetime import datetime
from .config import DATABASE_URL
from typing import List, Dict, Any

class DatabaseManager:
    """
    DatabaseManager 用于管理数据库连接、结构信息获取和SQL执行。
    支持获取表结构、字段、数据，清理和校验SQL等功能。
    """
    def __init__(self):
        """
        初始化DatabaseManager，创建数据库引擎并加载表结构信息。
        """
        self.engine: Engine = create_engine(DATABASE_URL)
        self.table_info = self._load_table_info()
    
    def _convert_value(self, value):
        """
        转换特殊类型的值（如datetime为字符串）。
        参数：
            value: 任意类型的值。
        返回：
            转换后的值。
        """
        if isinstance(value, datetime):
            return value.isoformat()
        return value
    
    def execute_query(self, query: str) -> list:
        """
        执行SQL查询并返回结果。
        参数：
            query (str): SQL语句。
        返回：
            list: 查询结果，每行为字典。
        """
        with self.engine.connect() as connection:
            result = connection.execute(text(query))
            # 获取列名
            columns = result.keys()
            # 将结果转换为字典列表，同时处理特殊类型
            return [
                {
                    col: self._convert_value(value)
                    for col, value in zip(columns, row)
                }
                for row in result
            ]
    
    def get_tables(self) -> list:
        """
        获取所有表名及其注释。
        返回：
            list: 每个元素为{'table_name', 'table_comment'}。
        """
        query = """
        SELECT 
            t.table_name,
            obj_description((quote_ident(t.table_schema) || '.' || quote_ident(t.table_name))::regclass, 'pg_class') as table_comment
        FROM information_schema.tables t
        WHERE t.table_schema = 'public' 
        ORDER BY t.table_name;
        """
        return self.execute_query(query)
    
    def get_table_columns(self, table_name: str) -> list:
        """
        获取指定表的所有列信息及其注释。
        参数：
            table_name (str): 表名。
        返回：
            list: 字段信息列表。
        """
        query = f"""
        SELECT 
            c.column_name,
            c.data_type,
            c.is_nullable,
            col_description((quote_ident(c.table_schema) || '.' || quote_ident(c.table_name))::regclass::oid, c.ordinal_position) as column_comment
        FROM information_schema.columns c
        WHERE c.table_name = '{table_name}'
        ORDER BY c.ordinal_position;
        """
        return self.execute_query(query)
    
    def get_table_data(self, table_name: str, limit: int = 10) -> list:
        """
        获取指定表的数据。
        参数：
            table_name (str): 表名。
            limit (int): 返回的最大行数，默认10。
        返回：
            list: 表数据（每行为字典）。
        """
        query = f"""
        SELECT *
        FROM {table_name}
        LIMIT {limit};
        """
        return self.execute_query(query)

    def clean_sql(self, sql: str) -> str:
        """
        清理SQL语句，移除代码块标记和think部分。
        参数：
            sql (str): 原始SQL语句。
        返回：
            str: 清理后的SQL语句。
        """
        # 移除可能的代码块标记
        sql = sql.replace('``sql', '').replace('```', '').strip()
        
        # 移除think部分
        if '<think>' in sql:
            # 找到think部分的开始和结束
            think_start = sql.find('<think>')
            think_end = sql.find('</think>')
            if think_end != -1:
                # 提取think部分后的SQL语句
                sql = sql[think_end + 8:].strip()
            else:
                # 如果没有找到结束标记，只保留think部分之后的内容
                sql = sql[think_start + 7:].strip()
        
        # 移除SQL语句前后的空白字符
        sql = sql.strip()
        
        # 如果SQL语句以分号结尾，移除分号
        if sql.endswith(';'):
            sql = sql[:-1]
            
        return sql
    
    def get_column_names(self, table_name: str) -> List[str]:
        """
        获取表的列名列表。
        参数：
            table_name (str): 表名。
        返回：
            list: 列名字符串列表。
        """
        if table_name in self.table_info:
            return [col['column_name'] for col in self.table_info[table_name]]
        return []
    
    def validate_and_fix_sql(self, sql: str) -> (bool, str):
        """
        验证并修复SQL语句。
        参数：
            sql (str): 原始SQL语句。
        返回：
            (bool, str): (验证是否通过, 验证通过后的SQL或失败原因)
        """
        # 只保留基本的SQL清理功能
        sql = self.clean_sql(sql)
        # 尝试解析和执行EXPLAIN，验证SQL语法和可执行性
        try:
            with self.engine.connect() as connection:
                # 只做语法和结构检查，不实际执行
                connection.execute(text(f"EXPLAIN {sql}"))
            return True, sql
        except Exception as e:
            return False, str(e)
    
    def _load_table_info(self) -> Dict[str, List[Dict[str, str]]]:
        """
        加载数据库表结构信息。
        返回：
            dict: {表名: {description, columns}}。
        """
        tables = self.get_tables()
        table_info = {}
        
        for table in tables:
            table_name = table['table_name']
            table_comment = table['table_comment'] or f"表 {table_name}"
            columns = self.get_table_columns(table_name)
            
            # 构建字段描述列表
            column_descriptions = []
            for col in columns:
                col_name = col['column_name']
                col_type = col['data_type']
                col_nullable = "可空" if col['is_nullable'] == 'YES' else "非空"
                col_comment = col['column_comment'] or f"字段 {col_name}"
                column_descriptions.append(f"{col_name} ({col_type}, {col_nullable}): {col_comment}")
            
            # 构建表描述
            table_info[table_name] = {
                'description': table_comment,
                'columns': column_descriptions
            }
        
        return table_info
    
    def generate_table_description(self) -> str:
        """
        生成数据库表结构的描述文本。
        返回：
            str: 数据库结构描述。
        """
        table_info = self._load_table_info()
        description = "数据库包含以下表：\n\n"
        
        for table_name, info in table_info.items():
            description += f"表名：{table_name}\n"
            description += f"描述：{info['description']}\n"
            description += "字段列表：\n"
            for col_desc in info['columns']:
                description += f"  - {col_desc}\n"
            description += "\n"
        
        return description 